from mozi_simu_sdk.mozi_server import MoziServer
import numpy as np
import threading
import time
import haversine


class MoZienv(object):

    def __init__(self, ip, port, platform=None, scenario_name=None, simulate_compression=5, duration_interval=15,
                 synchronous=True, app_mode=None, agent_key_event_file=None, request_id=None, side_name=None, **kwargs):
        # 墨子服务端IP
        self.server_ip = ip
        # 墨子服务端端口
        self.aiPort = port
        # 平台： windows 或 linux
        self.platform = platform
        # 想定名称
        self.scenario_name = scenario_name
        # app_mode:
        # 1--local windows windows模式
        # 2--linux mode    linux模式
        # 3--evaluate mode 比赛模式
        self.app_mode = app_mode
        # 推演档位 0-1倍速，1-2倍速，2-5倍速，3-15倍速，4-30倍速，5-60倍速，6-300倍速，7-900倍速，8-1800倍速
        self.simulate_compression = simulate_compression
        # 决策步长，单位秒
        self.duration_interval = duration_interval
        self.synchronous = synchronous

        # 用于比赛，智能体关键事件文件绝对路径，用于判定智能体是否正常运行
        self.agent_key_event_file = agent_key_event_file
        # 智能体身份标识
        self.request_id = request_id
        # 智能体扮演方，加载想定时用于清除扮演方任务
        self.side_name = side_name
        self.side = None
        # 决策
        self.step_count = 0
        self.mozi_server = None
        self.scenario = None
        self.situation = None
        self.last_score = 0
        self.last_DamageState = [0, 0, 0]
        self.last_distances = [0, 0, 0, 0]
        # 推进模式  True 同步 ,False 异步
        self.episode_limit = 1000
        self.n_agents = 17
        

        self.avail_actions = np.ones((17, 4))
        self.sum_DF = 4
        self.sum_YJ18_052d = 8
        self.sum_YJ100_055 = 16
        for i in range(8, 17):
            for j in range(2, 4):
                self.avail_actions[i][j] = 0
        self.action_mask = np.ones(17)
        self.state = [[0], [0], [0]]

    def step(self, actions):  
        global total_step
        if self.step_count == 0:
            self.mozi_server.run_simulate(self.app_mode)
            # zjy
            if self.app_mode == 3:
                task = threading.Thread(target=self.mozi_server.stream_send_and_recv, args=())
                task.setDaemon(True)  # 必须在start之前设置
                task.start()
        else:
            # if self.app_mode != 3:
            self.mozi_server.run_grpc_simulate(self.app_mode)
        self.step_count += 1
        self.situation = self.mozi_server.update_situation(self.scenario, self.app_mode)
        for k, si in self.scenario.get_sides().items():
            si.static_update()
        


        if self.is_done():  # 判断想定是否完成。
            terminated = True
        else:
            terminated = False
        info = {}
        info["episode_limit"] = False
        return reward, terminated, info

    def get_obs(self):
        """ Returns all agent observations in a list """
        return [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                ]

    def get_obs_agent(self, agent_id):
        """ Returns observation for agent_id """
        raise NotImplementedError

    def get_obs_size(self):
        """ Returns the shape of the observation """
        return 17

    def get_state(self):   
        aircrafts = [0]
        bships = [0]
        rships = [0]
        if self.scenario != None:
            for k, v in self.scenario.get_sides().items():
                if k == '7709784c-8be3-4823-8f9c-d69c5189f3b1':
                    for k1, v1 in v.aircrafts.items():
                        if is_in_poly([v1.dLongitude, v1.dLatitude], self.air_engagement_zoom):
                            aircrafts.append([v1.dLongitude, v1.dLatitude, 1])
                            aircrafts[0] = aircrafts[0] + 1
                    for k2, v2 in v.ships.items():
                        if is_in_poly([v2.dLongitude, v2.dLatitude], self.fleet_zoom):
                            rships.append([v2.dLongitude, v2.dLatitude, 1])
                            rships[0] = rships[0] + 1
                else:
                    for k1, v1 in v.aircrafts.items():
                        if is_in_poly([v1.dLongitude, v1.dLatitude], self.air_engagement_zoom):
                            aircrafts.append([v2.dLongitude, v2.dLatitude, 2])
                            aircrafts[0] = aircrafts[0] + 1
                    for k2, v2 in v.ships.items():
                        if is_in_poly([v2.dLongitude, v2.dLatitude], self.enemy_fleet_zoom):
                            bships.append([v2.dLongitude, v2.dLatitude, 2])
                            bships[0] = bships[0] + 1
                        if is_in_poly([v2.dLongitude, v2.dLatitude], self.fleet_zoom):
                            rships.append([v2.dLongitude, v2.dLatitude, 2])
                            rships[0] = rships[0] + 1
            self.state = list_together(aircrafts, rships, bships)
        else:
            self.state = [[aircrafts,
                           rships,
                           bships]]
        return self.state

    def get_state_size(self):
        return [1, 3, 25, 3]

    def get_avail_actions(self):
        return self.avail_actions

    def get_total_actions(self):
        return 4

    def reset(self):
        self.step_count = 0
        self.avail_actions = np.ones((17, 4))
        for i in range(8, 17):
            for j in range(2, 4):
                self.avail_actions[i][j] = 0
        self.create_scenario()
        self.mozi_server.set_simulate_compression(self.simulate_compression)
        if self.side_name:
            self.mozi_server.clean_mission_by_side_name(self.side_name)
        self.mozi_server.init_situation(self.scenario, self.app_mode)
        for k, side in self.scenario.get_sides().items():
            side.static_construct()

        self.init_last_dis()  # 初始化距离表
        return self.scenario

    def create_scenario(self):
        self.scenario = self.mozi_server.load_scenario()

    def close(self):
        pass

    def get_env_info(self):
        env_info = {"state_shape": self.get_state_size(),
                    "obs_shape": self.get_obs_size(),
                    "n_actions": self.get_total_actions(),
                    "n_agents": self.n_agents,
                    "episode_limit": self.episode_limit}
        return env_info

    def create_scenario(self):
 
        self.scenario = self.mozi_server.load_scenario()

    def connect_mozi_server(self, ip=None, port=None):
 
        if ip is None and port is None:
            self.mozi_server = MoziServer(self.server_ip, self.aiPort, self.platform, self.scenario_name,
                                          self.simulate_compression, self.synchronous, self.request_id,
                                          self.agent_key_event_file)
        elif ip is not None and port is not None:
            self.mozi_server = MoziServer(ip, str(port), self.platform, self.scenario_name,
                                          self.simulate_compression, self.synchronous, self.request_id,
                                          self.agent_key_event_file)
        time.sleep(4.0)

    def start(self, ip=None, port=None):
 
        if ip is None and port is None:
            self.connect_mozi_server()
        elif ip is not None and port is not None:
            self.connect_mozi_server(ip, port)
        else:
            raise ValueError('请正确配置墨子IP与端口！！！')

        self.mozi_server.set_run_mode(self.synchronous)
        self.mozi_server.set_decision_step_length(self.duration_interval)

    def is_done(self):
 
        response_dic = self.scenario.get_responses()
        for _, v in response_dic.items():
            if v.Type == 'EndOfDeduction':
                print('打印出标记：EndOfDeduction')
                if self.agent_key_event_file:
                    self.mozi_server.write_key_event_string_to_file('推演结束！')
                return True
        return False

    def get_reward(self):
 
        reward = (self.side.get_score() - self.last_score) / 595  # 计算得分奖励
        self.last_score = self.side.get_score()
        reward += len(self.side.ships) * 0.005
         

        return reward

    def init_last_dis(self):
 
        for i in range(4):
            aiicraft_loc = (self.side.aircrafts[self.aircrafts_guid[i + 9]].dLatitude,
                            self.side.aircrafts[self.aircrafts_guid[i + 9]].dLongitude)
            self.last_distances[i] = haversine.haversine(ship_loc, aiicraft_loc, unit=haversine.Unit.KILOMETERS)


def is_in_poly(p, poly):

    px, py = p
    is_in = False
    for i, corner in enumerate(poly):
        next_i = i + 1 if i + 1 < len(poly) else 0
        x1, y1 = corner
        x2, y2 = poly[next_i]
        if (x1 == px and y1 == py) or (x2 == px and y2 == py):  # if point is on vertex
            is_in = True
            break
        if min(y1, y2) < py <= max(y1, y2):  # find horizontal edges of polygon
            x = x1 + (py - y1) * (x2 - x1) / (y2 - y1)
            if x == px:  # if point is on edge
                is_in = True
                break
            elif x > px:  # if point is on left-side of line
                is_in = not is_in
    return is_in


def list_together(list1, list2, list3):
    for i in [list1, list2, list3]:
        for j in range(25 - len(i) + 1):
            i.append([0, 0, 0])
        i.pop(0)
    return [list1, list2, list3]
